

from sketchc import *
import numpy as np


window = Window3((1000, 800))

class Maze(Obj):
    def __init__(self, unit, l, count):
        self.unit = unit
        self.l = np.array(l, dtype="f4")
        self.count = count

        self.mode = moderngl.TRIANGLES
        self.t = Turtle()
        self.box = None

        self.prog =  {
            "vertex_shader":'''
#version 330 core
in vec3 aPos;
in vec3 aNormal;

uniform mat4 modelMat;
uniform mat4 viewMat;
uniform mat4 projMat;
//uniform mat4 normalMat;
uniform vec3[100] bbb;

  out vec4 vPos;
  out vec3 wNormal;     // world
  out vec3 vNormal;

void main() {
    vPos = modelMat * vec4(aPos, 1.0) + vec4(bbb[gl_InstanceID], 0);
    wNormal = mat3( modelMat) * aNormal;
    vNormal = mat3( viewMat * modelMat) * aNormal;
    //vNormal = mat3(transpose(inverse(modelMat))) * aNormal;
    //vNormal = mat3(normalMat) * aNormal;
    gl_Position = projMat * viewMat * vPos;
  }
            ''',
            "fragment_shader":'''
#version 330 core
out vec4 FragColor;

uniform vec3 lightPos;
uniform vec3 viewPos;
uniform vec3 objectColor;
uniform int get_normal;

in vec4 vPos;
in vec3 wNormal;
in vec3 vNormal;


void main()
{
    if (get_normal == 3){
        vec3 norm = normalize(vNormal);
        //FragColor = vec4(norm / 2.0 + vec3(0.5, 0.5, 0.5), 1.0);
        FragColor = vec4(norm / 2.0 + vec3(0.5, 0.5, 0.5), gl_FragCoord.z);
        return;
    }


    vec3 FragPos = vPos.xyz;
    float ambientStrength = 0.1;
    vec3 ambient = ambientStrength * objectColor;

    vec3 norm = normalize(wNormal);
    vec3 lightDir = normalize(lightPos - FragPos);
    float diff = max(dot(norm, lightDir) , 0.0);
    //vec3 diffuse = diff * vec3(0.5, 0.5, 0.5) ;
    vec3 diffuse = diff * objectColor ;
    diffuse = diffuse * 0.3 + 0.7;


    vec3 viewDir = normalize(viewPos - FragPos);
    float spec = 0.0;
    vec3 halfwayDir = normalize(lightDir + viewDir);
    spec = pow(max(dot(wNormal, halfwayDir), 0.0), 32.0);
    vec3 specular = spec * vec3(0.3);

    vec3 result = (ambient + diffuse + specular * 0.1);

    //vec3 red = vec3(1, 0, 0);

    FragColor = vec4(result, 1.0);
    //FragColor = vec4(result * 0 + objectColor, 1.0);
}
            '''}
            

        self.vao = None


    def __rrshift__(self, obj):
        raise NotImplemented()
    
    def u(self, ctx):
        assert isinstance(ctx, moderngl.Context)
        self.prog = ctx.program(**self.prog)
        self.vao= self.get_vao(ctx)
    
    def get_vao(self, ctx):
        vertices = np.array(self.unit.vertices, dtype="f4")
        normals = np.array(self.unit.normals, dtype="f4")
        # ~ print("::", vertices.shape)
        # ~ print("::", normals.shape)
        a, b, c = vertices[:,0], vertices[:,1], vertices[:,2]
        d, e, f = normals[:,0], normals[:,1], normals[:,2]
        x = np.dstack([a, b, c, d, e, f ])
        vbo = ctx.buffer(x.astype("f4"))
        return ctx.simple_vertex_array(self.prog, vbo, 'aPos', 'aNormal')
    
    def draw(self, camera):
        if self.vao is None:
            print("没准备好绘制", self)
            return

        self.prog['projMat'].write(camera.p())
        self.prog['modelMat'].write(self.m())
        self.prog['viewMat'].write(camera.v())
        self.prog['lightPos'].write(camera.pos)
        self.prog['viewPos'].write(camera.pos)
        self.prog['objectColor'].write(V3(0.5, 0.5, 0.5))
        self.prog['get_normal'].value = 0
        self.prog['bbb'].write(self.l)
        self.vao.render(self.mode, instances=self.count)
        # ~ self.vao.render(self.mode)
        

    def draw_normal(self, camera):
        if self.vao is None:
            print("没准备好绘制", self)
            return

        self.prog['modelMat'].write(self.m())
        self.prog['viewMat'].write(camera.v())
        self.prog['projMat'].write(camera.p())
        self.prog['lightPos'].write(camera.pos)
        self.prog['viewPos'].write(camera.pos)
        self.prog['objectColor'].write(V3(0.5, 0.5, 0.5))
        self.prog['get_normal'].value = 3
        # ~ self.vao.render(self.mode)
        self.prog['bbb'].write(self.l)
        self.vao.render(self.mode, instances=self.count)




# ~ box1 >> window

###随机部分
x, y = np.meshgrid(range(10), range(10))
z = np.zeros((10, 10))
r = np.sin(1 + x * 1.1 + y * 102.3) * 100
r =( r - np.floor(r)).flatten()
# ~ print(r)
p = np.dstack([x, z, -y]).reshape((100, 3)).astype("f4")
# ~ print(p)
p = p * 40
p1 = p[r < 0.5]
p2 = p[r > 0.5]
p11 = np.concatenate((p1, p2))
p22 = np.concatenate((p2, p1))
###随机部分结束

def madd(l, a):
    return [v + a for v in l]

def with_norm(l):
    return l, easy_norm(l)


def pspace(p1, p2, k):
    dp = p2 - p1
    space = np.linspace(0, 1, k)
    return [p1 + dp * k for k in space]

m = 2
n = 0.7
bottom = U3(0, 0, 0)
# ~ x = [U3(2, 1.3, 0), U3(2, 2, 0), U3(1.3, 2, 0), U3(-2, -1.3, 0), U3(-2, -2, 0), U3(-1.3, -2, 0)]
x = [U3(m, m-n, 0), U3(m, m, 0), U3(m-n, m, 0), U3(-m, -m+n, 0), U3(-m, -m, 0), U3(-m+n, -m, 0)]
x1 = [v + bottom for v in x]
top = U3(0, 0, 5)
x2 = [v + top for v in x]
vertices = []
vertices += make_ring([bottom], x1)
vertices += make_ring(x1, x2)
vertices += make_ring(x2, [top])

corner1 = a, b, c = U3(-2, 2, 0), U3(-2, 1.3, 0), U3(-1.3, 2, 0)
mi = (a + b + c) / 3
vertices += make_ring([mi + bottom], madd([a, b, c], bottom))
vertices += make_ring(madd([a, b, c], bottom), madd([a, b, c], top))
vertices += make_ring(madd([a, b, c], top), [mi + top])

corner2 = a, b, c = U3(2, -2, 0), U3(2, -1.3, 0), U3(1.3, -2, 0)
mi = (a + b + c) / 3
vertices += make_ring([mi + bottom], madd([a, b, c], bottom))
vertices += make_ring(madd([a, b, c], bottom), madd([a, b, c], top))
vertices += make_ring(madd([a, b, c], top), [mi + top])

vertices = [v * 10 for v in vertices]
o1 = Obj()
(vertices, easy_norm(vertices)) >> o1
# ~ o1.u(window.ctx)
# ~ o1 >> window

anchor1 = U3(m-n, m, 0) * 10
anchor2 = U3(-m, -m+n, 0) * 10

a1 = 0.3
h = 50
linev = make_block(U3(a1, a1, h))

for p in pspace(anchor1, anchor2, 33):
    turtle = Turtle(p, vrot(U3.up, pi/4))
    turtle.pile(linev, anchor=U3(0, 0, -1))
    linev.as_unit() >> o1

anchor3, anchor4 = anchor1, anchor2

anchor2 = U3(-m+n, -m, 0) * 10
anchor1 = U3(m, m-n, 0) * 10
    
for p in pspace(anchor1, anchor2, 33):
    turtle = Turtle(p, vrot(U3.up, pi/4))
    turtle.pile(linev, anchor=U3(0, 0, -1))
    linev.as_unit() >> o1


dh = 2.5

hori = Obj()
turtle1 = Turtle(anchor3+U3(0, 0,dh -  a1/2), vrot(U3.up, pi/4))    
turtle2 = Turtle(anchor4+U3(0, 0, dh - a1/2), vrot(U3.up, pi/4))    
ll = [U3(0, a1/2, a1 / 2), U3(0, a1/2, -a1 / 2), U3(0, -a1/2, -a1 / 2), U3(0, -a1/2, a1 / 2)]    
ll1 = turtle1.postl(ll)
ll2 = turtle2.postl(ll)
with_norm(make_ring(ll1, ll2)) >> hori

turtle1 = Turtle(anchor1+U3(0, 0,dh -  a1/2), vrot(U3.up, pi/4))    
turtle2 = Turtle(anchor2+U3(0, 0, dh - a1/2), vrot(U3.up, pi/4))    
ll1 = turtle1.postl(ll)
ll2 = turtle2.postl(ll)
with_norm(make_ring(ll1, ll2)) >> hori

turtle1 = Turtle(corner1[2] * 10+U3(0, 0,dh -  a1/2), vrot(U3.up, pi/4))    
turtle2 = Turtle(corner1[1] * 10 +U3(0, 0, dh - a1/2), vrot(U3.up, pi/4))    
ll1 = turtle1.postl(ll)
ll2 = turtle2.postl(ll)
with_norm(make_ring(ll1, ll2)) >> hori

turtle1 = Turtle(corner2[1] * 10+U3(0, 0,dh -  a1/2), vrot(U3.up, pi/4))    
turtle2 = Turtle(corner2[2] * 10+U3(0, 0, dh - a1/2), vrot(U3.up, pi/4))    
ll1 = turtle1.postl(ll)
ll2 = turtle2.postl(ll)
with_norm(make_ring(ll1, ll2)) >> hori



for i in range(20):
    hori.as_unit() >> o1
    hori.t.walk(U3(0,0,dh))

##

anchor1 = U3(-m, m-n, 0) * 10
anchor2 = U3(-m+n, m, 0) * 10
for p in pspace(anchor1, anchor2, 7):
    turtle = Turtle(p, vrot(U3.up, pi/4))
    turtle.pile(linev, anchor=U3(0, 0, -1))
    linev.as_unit() >> o1

    
anchor1 = U3(m, -m+n, 0) * 10
anchor2 = U3(m-n, -m, 0) * 10
for p in pspace(anchor1, anchor2, 7):
    turtle = Turtle(p, vrot(U3.up, pi/4))
    turtle.pile(linev, anchor=U3(0, 0, -1))
    linev.as_unit() >> o1



################################### o1

x = [U3(-m+n, m, 0), U3(-m, m, 0), U3(-m, m-n, 0) , U3(m-n, -m, 0) , U3(m, -m, 0), U3(m, -m+n, 0)]
x1 = [v + bottom for v in x]
x2 = [v + top for v in x]
vertices = []
vertices += make_ring([bottom], x1)
vertices += make_ring(x1, x2)
vertices += make_ring(x2, [top])

a, b, c = U3(1.3, 2, 0), U3(2, 1.3, 0), U3(2, 2, 0)
mi = (a + b + c) / 3
vertices += make_ring([mi + bottom], madd([a, b, c], bottom))
vertices += make_ring(madd([a, b, c], bottom), madd([a, b, c], top))
vertices += make_ring(madd([a, b, c], top), [mi + top])

a, b, c = U3(-1.3, -2, 0), U3(-2, -1.3, 0), U3(-2, -2, 0)
mi = (a + b + c) / 3
vertices += make_ring([mi + bottom], madd([a, b, c], bottom))
vertices += make_ring(madd([a, b, c], bottom), madd([a, b, c], top))
vertices += make_ring(madd([a, b, c], top), [mi + top])

vertices = [v * 10 for v in vertices]

o2 = Obj()
(vertices, easy_norm(vertices)) >> o2
# ~ o2.u(window.ctx)
# ~ o2 >> window


anchor1 = U3(-m, m-n, 0) * 10
anchor2 = U3(m-n, -m, 0) * 10

a1 = 0.3
h = 50
linev = make_block(U3(a1, a1, h))

for p in pspace(anchor1, anchor2, 33):
    turtle = Turtle(p, vrot(U3.up, pi/4))
    turtle.pile(linev, anchor=U3(0, 0, -1))
    linev.as_unit() >> o2


anchor2 = U3(-m+n, m, 0) * 10
anchor1 = U3(m, -m+n, 0) * 10
    
for p in pspace(anchor1, anchor2, 33):
    turtle = Turtle(p, vrot(U3.up, pi/4))
    turtle.pile(linev, anchor=U3(0, 0, -1))
    linev.as_unit() >> o2


anchor1 = U3(-m+n, -m, 0) * 10
anchor2 = U3(-m, -m+n, 0) * 10
for p in pspace(anchor1, anchor2, 7):
    turtle = Turtle(p, vrot(U3.up, pi/4))
    turtle.pile(linev, anchor=U3(0, 0, -1))
    linev.as_unit() >> o2

    
anchor1 = U3(m, m-n, 0) * 10
anchor2 = U3(m-n, m, 0) * 10
for p in pspace(anchor1, anchor2, 7):
    turtle = Turtle(p, vrot(U3.up, pi/4))
    turtle.pile(linev, anchor=U3(0, 0, -1))
    linev.as_unit() >> o2



hori.vertices = [V3(-u[0], u[1], u[2]) for u in hori.vertices][::-1]
hori.normals = [V3(-u[0], u[1], u[2]) for u in hori.normals][::-1]
hori.t.pos = U3(0,0,0)
for i in range(20):
    hori.as_unit() >> o2
    hori.t.walk(U3(0,0,dh))

###################################o2 divide

# ~ box1 = make_block(U3(0.3, 0.3, 0.3))
m1 = Maze(o1, p11, len(p1))
m1.u(window.ctx)
m1 >> window

# ~ box1 = make_block(U3(0.6, 0.6, 0.6))
m2 = Maze(o2, p22, len(p2))
m2.u(window.ctx)
m2 >> window


camera = Camera(U3(10, -5, 50))
camera >> window        # 识别为camera则自动作为主camera

walker = Walker(camera)
walker >> window


clock = pygame.time.Clock()
while True:
    dt = clock.tick(60)
    window.update(dt)
    for e in pygame.event.get():
        if e.type == pygame.QUIT:
            exit()
        window.dispatch(e)
    window.render()   # screen 操作需要写在window的槽里
    pygame.display.flip()

